#ifndef OPTIMIZER_H
#define OPTIMIZER_H

#include <queue>
#include <unordered_map>
#include <unordered_set>

#include "scan_matching/occupied_grid_cost_function_2d.h"
#include "scan_matching/pose_observe.h"
#include "scan_matching/state.h"

using namespace std;
using namespace Eigen;
using namespace cartographer; 
using namespace scan_matching; 

static const double PI = 3.14159265358979323846;

class RobustKernel {
public:
    virtual ~RobustKernel() = default;
    virtual tuple<MatrixXd, VectorXd> robustify(const VectorXd& e) const = 0;
    virtual double weight(double squared_error) const = 0;
};

class HuberKernel : public RobustKernel {
public:
    HuberKernel(double delta = 1.0) : delta(delta), delta_sq(delta * delta) {}
    
    tuple<MatrixXd, VectorXd> robustify(const VectorXd& e) const override {
        double e_sq = e.squaredNorm();
        double w = weight(e_sq);
        VectorXd weighted_e = w * e;
        MatrixXd weighted_H = w * MatrixXd::Identity(e.size(), e.size());
        return make_tuple(weighted_H, weighted_e);
    }
    
    double weight(double squared_error) const override {
        return squared_error > delta_sq ? delta / sqrt(squared_error) : 1.0;
    }

private:
    double delta;
    double delta_sq;
};

struct PoseGraphEdge {
    int id1, id2;
    State measurement;
    MatrixXd information;
    bool is_outlier = false;
    
    PoseGraphEdge(int id1, int id2, const State& meas, const MatrixXd& info)
        : id1(id1), id2(id2), measurement(meas), information(info) {}
};

struct MarginalizationFactor {
    MatrixXd H_prior;
    VectorXd b_prior;
    unordered_set<int> related_vertices;
};

class Optimizer {
public:
	Optimizer(std::shared_ptr<OccupiedGridCostFunction2D> grid_cost_function_ptr);
	void Optimize(const int& max_iterations = 100);
	int AddVertex(const State& state, const bool& is_fixed = false);
	void AddEdge(const PoseGraphEdge& edge);

    void UpdateState(State& state, const int id);


    void PrintPoses() const {
        cout << "Current poses:" << endl;
        for (const auto& v : vertices_) {
            cout << "Vertex " << v.first << ": (" << v.second.p[0] << ", " 
                 << v.second.p[1] << ", " << SO2(v.second.R).log() << ")" << endl;
        }
    }

private:
	void MarginalizeOldestVertex();
	double ComputeTotalError();
	void BuildSystem();
	void UpdatePoses(const VectorXd& dx);

	std::shared_ptr<OccupiedGridCostFunction2D> grid_cost_function_ptr_;

	int window_size_ = 1;
    int next_vertex_id_ = 0;
    queue<int> vertex_queue_;
    unordered_map<int, State> vertices_;
    unordered_map<int, bool> is_fixed_;
    vector<PoseGraphEdge> edges_;
    unordered_set<size_t> active_edges_;
    unordered_map<int, unordered_set<size_t>> vertex_to_edges_;
    unordered_map<int, MarginalizationFactor> marginalization_factors_;
    shared_ptr<RobustKernel> robust_kernel_;
    
    MatrixXd H_;
    VectorXd b_;
    unordered_map<int, int> var_indices_;

    bool is_solve_success_;
};

#endif
